from solvers.solverbase import SolverBase
from solvers.dg.dgbase import DGBase
from dolfin import *

class CrankNicolson(SolverBase,DGBase):
    'Discontinuous Galerkin Crank-Nicolson method.'
    def __init__(self,p=1,inflowBC=0):
        SolverBase.__init__(self,p)
        DGBase.__init__(self,inflowBC)

    def solve(self,problem):
        mesh = problem.domain

        p_degree = self.p
        V = FunctionSpace(mesh,'DG',p_degree)
        VV = VectorFunctionSpace(mesh,'DG',p_degree)

        phi_ = interpolate(problem.phi_,V)
        phi = TrialFunction(V)
        varphi = TestFunction(V)

        t = problem.t           
        dt = Constant(SolverBase._get_time_step(self,problem,VV,t))
        T = problem.T

        v = SolverBase._get_velocity(self,problem,VV,t)

        
        self.update(problem,phi_,t,float(dt),0)  # store the solution at time 0
        n = FacetNormal(mesh)
       
        if not self.inflowBC: # no inflow boundary conditions
            a = phi*varphi*dx - 0.5*dt*inner(v,nabla_grad(varphi))*phi*dx\
                + 0.5*dt('+')*dot(v('+'),jump(varphi,n)*avg(phi))*dS\
                + 0.5*dt('+')*0.5*abs(dot(v('+'),n('+')))*dot(jump(phi,n),jump(varphi,n))*dS\
                + 0.5*dt*phi*varphi*dot(v,n)*ds

            l = phi_*varphi*dx + 0.5*dt*inner(v,nabla_grad(varphi))*phi_*dx\
                - 0.5*dt('+')*dot(v('+'),jump(varphi,n)*avg(phi_))*dS\
                - 0.5*dt('+')*0.5*abs(dot(v('+'),n('+')))*dot(jump(phi_,n),jump(varphi,n))*dS\
                - 0.5*dt*phi_*varphi*dot(v,n)*ds
        else:
            (ib,ib_value,bParts) = problem.get_inflowBCS()  # ib_value has bc conditions on the inflow part
            g = Function(V)                                 # inflow bcs for the weak form
 
            a = phi*varphi*dx + 0.5*dt*inner(v,nabla_grad(phi))*varphi*dx\
                - 0.5*dt('+')*dot(v('+'),jump(phi,n)*avg(varphi))*dS\
                + 0.5*dt('+')*0.5*abs(dot(v('+'),n('+')))*dot(jump(phi,n),jump(varphi,n))*dS\
                + 0.5*dt*phi*varphi*0.5*(abs(dot(v,n))-dot(v,n))*ds

            l = phi_*varphi*dx - 0.5*dt*inner(v,nabla_grad(phi_))*varphi*dx\
                + 0.5*dt('+')*dot(v('+'),jump(phi_,n)*avg(varphi))*dS\
                - 0.5*dt('+')*0.5*abs(dot(v('+'),n('+')))*dot(jump(phi_,n),jump(varphi,n))*dS\
                - 0.5*dt*(phi_-g)*varphi*0.5*(abs(dot(v,n))-dot(v,n))*ds\
                + 0.5*dt*g*varphi*0.5*(abs(dot(v,n))-dot(v,n))*ds

        A = assemble(a)     
        phi = Function(V)
   
        t = 0
        count = 0

        while t <= T:
            t += float(dt)
            count += 1
           
            if self.inflowBC:
                g.vector()[:] = interpolate(ib_value(t),V).vector().array() # update inflow values
         
            L = assemble(l) 

            solve(A,phi.vector(),L)
            phi_.assign(phi)
            
            self.update(problem,phi_,t,float(dt),count)

        problem.save_data(self.saveDir,str(self.CFL))

